"""
    lens shading correction

    author: wxz
    date: 
    github: https://github.com/xinzwang
"""

from isp.blocks.utils import *


def lens_shading_correction(img, flat_max=1.0, flat_min=0.9):
    height, width = img.shape

    dark = np.zeros_like(img, dtype=np.float32)
    flat = np.zeros_like(img, dtype=np.float32)

    center = [height / 2, width / 2]
    max_distance = distance(center, [height, width])

    for i in range(height):
        for j in range(width):
            dist = (max_distance - distance(center, [i, j])) / max_distance
            flat[i, j] = flat_min + dist * (flat_max - flat_min)

    t = flat - dark
    data = np.average(t) * np.divide((img - dark), t)

    np.clip(data, 0, 1)
    return data
